package vulnerabilityreport_test

import (
	"sort"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/aquasecurity/trivy-operator/pkg/apis/aquasecurity/v1alpha1"
	"github.com/aquasecurity/trivy-operator/pkg/vulnerabilityreport"
)

func TestBySeverity(t *testing.T) {
	items := []v1alpha1.Vulnerability{
		{VulnerabilityID: "CVE-0000-0001", Severity: v1alpha1.SeverityLow},
		{VulnerabilityID: "CVE-0000-0002", Severity: v1alpha1.SeverityMedium},
		{VulnerabilityID: "CVE-0000-0003", Severity: v1alpha1.SeverityMedium},
		{VulnerabilityID: "CVE-0000-0004", Severity: v1alpha1.SeverityCritical},
		{VulnerabilityID: "CVE-0000-0005", Severity: v1alpha1.SeverityHigh},
		{VulnerabilityID: "CVE-0000-0006", Severity: v1alpha1.SeverityHigh},
	}

	sort.Stable(vulnerabilityreport.BySeverity{Vulnerabilities: items})

	assert.Equal(t, []v1alpha1.Vulnerability{
		{VulnerabilityID: "CVE-0000-0004", Severity: v1alpha1.SeverityCritical},
		{VulnerabilityID: "CVE-0000-0005", Severity: v1alpha1.SeverityHigh},
		{VulnerabilityID: "CVE-0000-0006", Severity: v1alpha1.SeverityHigh},
		{VulnerabilityID: "CVE-0000-0002", Severity: v1alpha1.SeverityMedium},
		{VulnerabilityID: "CVE-0000-0003", Severity: v1alpha1.SeverityMedium},
		{VulnerabilityID: "CVE-0000-0001", Severity: v1alpha1.SeverityLow},
	}, items)

}

func TestOrderedBy(t *testing.T) {
	reports := []v1alpha1.VulnerabilityReport{
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "nginx",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 3,
					HighCount:     12,
					MediumCount:   58,
					LowCount:      7,
					UnknownCount:  1,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "wordpress",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 7,
					HighCount:     12,
					MediumCount:   1,
					LowCount:      20,
					UnknownCount:  2,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "mysql",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 3,
					HighCount:     2,
					MediumCount:   1,
					LowCount:      7,
					UnknownCount:  3,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "postgresql",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 89,
					HighCount:     0,
					MediumCount:   1,
					LowCount:      3,
					UnknownCount:  2,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "jenkins",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 0,
					HighCount:     0,
					MediumCount:   1,
					LowCount:      7,
					UnknownCount:  3,
				},
			},
		},
	}
	vulnerabilityreport.OrderedBy(vulnerabilityreport.SummaryCount...).SortDesc(reports)
	assert.Equal(t, []v1alpha1.VulnerabilityReport{
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "postgresql",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 89,
					HighCount:     0,
					MediumCount:   1,
					LowCount:      3,
					UnknownCount:  2,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "wordpress",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 7,
					HighCount:     12,
					MediumCount:   1,
					LowCount:      20,
					UnknownCount:  2,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "nginx",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 3,
					HighCount:     12,
					MediumCount:   58,
					LowCount:      7,
					UnknownCount:  1,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "mysql",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 3,
					HighCount:     2,
					MediumCount:   1,
					LowCount:      7,
					UnknownCount:  3,
				},
			},
		},
		{
			Report: v1alpha1.VulnerabilityReportData{
				Artifact: v1alpha1.Artifact{
					Repository: "jenkins",
				},
				Summary: v1alpha1.VulnerabilitySummary{
					CriticalCount: 0,
					HighCount:     0,
					MediumCount:   1,
					LowCount:      7,
					UnknownCount:  3,
				},
			},
		},
	}, reports)
}
